"""
Generic graph plotting functionality.

"""

import matplotlib.pyplot as plt
import numpy as np

class Plotter(object):
    """Class providing generic graph plotting methods."""

    markers = ['x', '+', 'o', 'v', 'd', '<', 's', '8']
    #: Available scatter plot markers for matplotlib
    marker_sizes = dict([ (key, 49) for key in markers ])
    #: Size to use for each marker (default = 49, except 'o' whose default is 9)
    marker_sizes['o'] = 9
    marker_ind = 0

    artists = []
    index_artists = []
    index_heights = []
    index_labels = []

    x_max = 0
    y_max = 0

    def plot_box(self, xvals, yvals, label=None, widths=None, *args, **kwargs):
        """Given one or more x-value, will plot all given y-values as a
        box-and-whisker plot.

        This is designed to be called once per dataset, so that each dataset
        has the same colour.

        See http://matplotlib.sourceforge.net/api/pyplot_api.html#matplotlib.pyplot.boxplot
        for more information"""
        if widths is None:
            widths = self.increment * 0.3
        mappings = self.axis.boxplot(yvals, *args, positions=xvals,
                                     widths=widths, **kwargs)
        for mapping in mappings:
            plt.setp(mappings[mapping], color=self.colours[self.col_ind])
        # mappings.keys() = ['medians', 'fliers', 'whiskers', 'boxes', 'caps']
        plt.setp(mappings['whiskers'], linestyle='dotted')
        self.col_ind += 1
        return mappings

    def plot_scatter(self, xvals, yvals,
                     marker=None, label=None, *args, **kwargs):
        """
        :param xvals: Passed directly to plt.scatter(), so must be compatible
                      with matplotlib.

        :param yvals: Can optionally be a dictionary of vectors, in which case
                      all the keys must be present in :data:`xvals`. If not a
                      dictionary, assume that it is a sequence with same length
                      as xvals.

        :param marker: Can be given a number indexing which of self.markers to
                       use, or a string to pass directly to matplotlib. By
                       default, rotate through markers, using whatever marker is
                       currently selected with: `self.markers[self.marker_ind]`

        label can optionally be used to indicate the label in the legend.

        Optional :data:`args` and :data:`kwargs` are passed directly to
        :func:`matplotlib.pyplot.scatter()`
        """
        all_xvals = []
        all_yvals = []
        if marker is None:
            marker = self.markers[self._marker_ind]
        elif type(marker) is int:
            marker = self.markers[marker]
        if type( yvals ) is dict:
            for n in xrange( len( yvals[xvals[0]] ) ):
                all_xvals.append( xvals )
                all_yvals.append( [ yvals[x][n] for x in xvals ] )
        else:
            all_xvals = xvals
            all_yvals = yvals
        axis = getattr(self, 'axis{0}'.format(self.cur_axis))
        try:
            # tests.
            mark_size = self.marker_sizes[marker]
            col = self.colours[self.col_ind]
        except IndexError:
            print('\n', len(self.marker_sizes), marker)
            print(len(self.colours), self.col_ind)
            raise
        else:
            # was just testing..
            del(mark_size)
            del(col)

        return axis.scatter(all_xvals, all_yvals, marker=marker,
                            s=self.marker_sizes[marker],
                            color=self.colours[self.col_ind],
                            label=label, *args, **kwargs)

    def plot_scatter_errors(self, xvals, yvals , mid=None, xerr=None,
                            yerr=None, label=None, *args, **kwargs ):
        """Plot xvals against yvals with error bars. Uses
        :data:`self.plot_scatter` to draw the points, but adds error bars.

        :param mid: This is the function which calculates the centre of
                    crosses. By default, use :func:`numpy.mean`.

        :param xerr: By default, don't draw x error bars. Otherwise, pass a
                     function reference which will return the error values when
                     passed xvals.

        :param yerr: By default, draw 95% Confidence Intervals. i.e.
                     :math:`\pm 1.96 *  \frac{(\STD. dev.}{\sqrt{replicates}}`
        """
        if mid is None:
            mid = np.mean
        mean_y_vals = [ mid(yval) for yval in yvals ]
        collection = self.plot_scatter(xvals, mean_y_vals,
                                       label=label, *args, **kwargs )
        xerrors = None
        if xerr is not None:
            xerrors = xerr( xvals )
        if yerr is None:
            yerrors = [ 1.96 * (np.std(yval) / np.sqrt(self.replicates)) \
                        for yval in yvals ]
            yerrors = [ np.std(yval) for yval in yvals ]
        else:
            yerrors = yerr(yvals)
        if label is not None:
            print('{0}\t'.format(label)),
        print('\t'.join(['{0:.3f} +- {1:.3f}'.format(mean_y_vals[i], yerrors[i])
                        for i in xrange(len(yvals))]))
        print('')
        y2_max = max(mean_y_vals) + max(yerrors)
        axis = getattr(self, 'axis{0}'.format(self.axis))

        (plotline, caplines, barlinecols) = \
            axis.errorbar(xvals, mean_y_vals, c=self.colours[self.col_ind],
                          xerr=xerrors, yerr=yerrors, linestyle='None',
                          label=label, *args, **kwargs)

        plt.setp(barlinecols, linestyles='dotted')
        if axis.get_ylim()[1] < y2_max:
            axis.set_ylim( 0 , int(y2_max) + 1 )
        return collection

## These two functions from:-
## http://matplotlib.sourceforge.net/examples/pylab_examples/multiple_yaxis_with_spines.html
def make_patch_spines_invisible(ax):
    ax.set_frame_on(True)
    ax.patch.set_visible(False)
    for sp in ax.spines.itervalues():
        sp.set_visible(False)

def make_spine_invisible(ax, direction):
    if direction in ["right", "left"]:
        ax.yaxis.set_ticks_position(direction)
        ax.yaxis.set_label_position(direction)
        anti = 'right' if direction == 'left' else 'left'
    elif direction in ["top", "bottom"]:
        ax.xaxis.set_ticks_position(direction)
        ax.xaxis.set_label_position(direction)
        anti = 'bottom' if direction == 'top' else 'top'
    else:
        raise ValueError("Unknown Direction : %s" % (direction,))
    ax.spines[direction].set_visible(True)
    ax.spines[anti].set_visible( False )
